#!/usr/bin/env python
'''
Given the PerM's output files of AceView's gene mapping
this program get all mappings of the isoforms related to the gene  
'''
__author__ = "Yangho Chen"
__email__ = "yanghoch@usc.edu"
__modification__ = "Created on Aug 16, 2009"

import sys
import os
import getopt
from PerM_output_format import *
from common import output_lines

def extract_gene_name ( isoform_name, prefix_delimiter, suffix_re_delimiter ):
    """ return string which exclude the substring 
    before prefix_delimiter and the after suffix_delimiter.
    If the string contains no suffix and prefix, return the original string """
    if ( isoform_name.lower().find( "ch" ) == 0 ):
        return( isoform_name ) # exception for chrM
    isoform_name = isoform_name.replace( prefix_delimiter, "" )
    end_index = isoform_name.rfind( suffix_re_delimiter )
    if( end_index >= 0 ):
        isoform_name = isoform_name[0:end_index]
    # else:        
    #    print "null gene_name in " + str( isoform_name )
    return( isoform_name )

def isAceViewIsoformName( isoform_name ):
    if ( "MRNA_" in isoform_name ):
        if isoform_name.index( "MRNA_" ) == 0:
            return True
    return False

def extract_AceView_gene_name( isoform_name ):
    ''' replace prefix of MRNA_ and everything after . at the end,
    which indicates the isoform id in AceView. 
    Return the original arg if it is not in the AceView's isoform format.'''
    tmp_isoform_name = isoform_name.replace( ':', '_' )
    tmp_isoform_name = tmp_isoform_name.replace( '-', '_' )
    if ( not isAceViewIsoformName( isoform_name ) ):
        return( isoform_name )
    else:
        return( extract_gene_name ( tmp_isoform_name, "MRNA_", "." ) )

def get_mapping4aGene( mapping_filename, target_gene_name ):
    """ 
    Given a mapping file and the gene name, get all the mapping 
    associate with the gene.
    """
    ofile = open( target_gene_name + ".mapping", mode = "a" )
    ifile = open( mapping_filename, mode = "r" )

    mappingNo = 0;
    while 1:
        line = ifile.readline()
        if not line:
            print "Total " + str( mappingNo ) + " mappings related to the gene";
            break;
        mapping = PerMMappingLineInfo ( line )
        gene_Name = extract_AceView_gene_name( mapping.isoform_name )
        if( gene_Name == target_gene_name ):
            ofile.write( line )
            mappingNo = mappingNo + 1;
    ifile.close()
    ofile.close()

def exclude_reads_mapped2multiple_genes( mapping_filename ):
    """ 
    Given a mapping file from PerM, exclude reads that mapped 
    to multiple AceView genes but allow reads mapped to multiple isoforms.
    """
    ifile = open( mapping_filename, mode = "r" )
    uniqueMappingExtName = "_unique.mapping"
    uniqueMappingFilename = \
        mapping_filename.replace( ".mapping", uniqueMappingExtName )
    ofile = open( uniqueMappingFilename, mode = "w" )
    readName = ""
    gene = ""
    mappings = set() # set to collect junction info per contig
    uniqueMapped = True
    uniqueMappedReadNo = 0;
    while 1:
        line = ifile.readline()
        if ( not line ) or line == "\n":
            #output the last 
            output_lines( ofile, mappings )
            print "Total " + str( uniqueMappedReadNo ) + " unique mapped reads";
            break;
        # TODO be careful if line is not in a correct format
        if( len( line ) == 0 ):
            print "Null line"

        mapping = PerMMappingLineInfo ( line )
        mapping.gene = extract_AceView_gene_name( mapping.isoform_name )
        if ( mapping.read_name != readName ):
            readName = mapping.read_name # new read
            gene = mapping.gene
            if ( uniqueMapped == True ):
                output_lines( ofile, mappings )
                uniqueMappedReadNo += 1
            mappings.clear()
            mappings.add( line )
            uniqueMapped = True
        else : # if the line is for different read     
            if( mapping.gene == gene and uniqueMapped ):
                mappings.add( line )
            else:
                uniqueMapped = False
    ifile.close()
def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print __doc__
    print "argv[1] is the mapping file name or the fast a file."
    print "argv[2] is the AceView Gene name."

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error, msg:
        print msg
        sys.exit( 2 );
    #process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            print __doc__
            sys.exit();

    if( len( args ) == 1 ):
        ext_name = os.path.splitext( args[0] )
        if( ext_name[1] == '.mapping' ):
            exclude_reads_mapped2multiple_genes( args[0] )
    elif( len( args ) == 2 ):
        max_coverage_of_ocean = int( args[1] )
        get_mapping4aGene( args[0], max_coverage_of_ocean )
    else:
        print "\nThe first parameter is the mapping file name\n"
        print "\nThe second parameter is the AceView Gene name\n"

if __name__ == '__main__':
    main()
